package com.wang.tool.okhttp;

import com.alibaba.fastjson.JSON;
import com.wang.tool.monitor.core.MonitorInfo;
import com.wang.tool.monitor.template.AbsMonitorInvoker;
import com.wang.tool.trace.TraceContext;
import com.wang.tool.trace.TraceHolder;
import com.wang.tool.trace.TraceTemplate;
import okhttp3.Interceptor;
import okhttp3.Request;
import okhttp3.Response;

import java.io.IOException;
import java.util.function.BiConsumer;

public class OkHttpTraceIdInterceptor implements Interceptor {

    @Override
    public Response intercept(Chain chain) throws IOException {
        // 获取原始请求
        Request originalRequest = chain.request();
        Request.Builder builder = originalRequest.newBuilder();
        return TraceTemplate.monitorWithNext("okhttp_intercept", (originalRequest.isHttps()?"https_":"http_")+originalRequest.method(),originalRequest.url().host()+':'+originalRequest.url().port()+ originalRequest.url().uri().getPath(), new AbsMonitorInvoker<Response>() {
            @Override
            public Response internalCall() throws Throwable {
                // 构建带有 traceId 的新请求
                return chain.proceed(builder.build());
            }
        }, new BiConsumer<TraceContext, MonitorInfo<Response>>() {

            @Override
            public void accept(TraceContext context, MonitorInfo<Response> monitorInfo) {
                if (context == null) {
                    return;
                }
                builder.addHeader(TraceHolder.TRACE_ID_HEADER, context.getTraceId())
                        .addHeader(TraceHolder.RPC_ID_HEADER, context.getRpcId());
                if (context.getAttributes() != null) {
                    builder.addHeader(TraceHolder.ATTRIBUTES_HEADER, JSON.toJSONString(context.getAttributes()));
                }
            }
        });
    }
}